/* 算法实现模板 */
import AlgorithmBase from "./AlgorithmBase";
import React from "react";
import ReactDom from 'react-dom';
import { loadDataset, txtToMatrix , mdRender} from "../../../common/Common";
import {guass,transpose,multiply} from "../../../common/EmFuncs"
import MatView2D, {createRandomPointStyleDict} from "../../../visual/MatView2D";
import Mat from "../../../common/Mat";

import 'katex/dist/katex.min.css';
import TDom from "../../../common/TDom";
import TCommon from "../../../common/TCommon";
import TMatrix from "../../../common/TMatrix";
import { showMathStringInTable } from "../../../visual/Table2D";

class EMForEoPOfGMM extends AlgorithmBase{
  constructor(){
    super('EMForEoPOfGMM');

    this.descriptionUrl = `${this.homepage}/data/slm/ch9/9.2/README.md`;
    this.datasetUrl     = `${this.homepage}/data/slm/ch9/9.2/dataset.txt`;

  }

  init(){ 
    this.initTrainDS();

    this.reset();
  }

  reset(){
    this.dataset = undefined;
    this.trainDS = undefined;

    this.clear();

    this.initDescription();
    this.updateModel();
  }
  
  // 模型UI
  initModelUI(){
    

    this.modelDiv = this.addDiv(this.domElementList.model);
  }

  // 渲染模型到页面
  updateModel(){
    var w = new Mat.zeros([1, 2]);
    var b = 0.0;

    const trainDS = this.trainDS;
    let X = new Mat.Mat(trainDS.Mat);
    var x=[]
    for(var i=0;i<X.length;i++){
      x.push(X[i])
    }
    this.x = x
    let k = 2, N = x.length, D = x[0].length

    //pi初值已设定
    var pi = []
    for(var i=0;i<k;i++) pi.push(1/k)
    var lpi = pi.concat()

    //mu初始值已设定
    var mu = [],lmu = []
    for(var i=0;i<k;i++){
        mu.push([])
        mu[i] = x[Math.floor(Math.random()*((N-1)/2)+i*N/2)]
        lmu.push([])
        lmu[i] = mu[i].concat()
    } 

    //求各维度均值
    var avg = []
    for(var i=0;i<D;i++) avg.push(0)
    for(var i=0;i<N;i++){
        for(var j=0;j<D;j++) avg[j]+=x[i][j]
    }
    for(var i=0;i<D;i++) avg[i] = (avg[i]/N)
    var x_ = []
    for(var i=0;i<N;i++){
        x_.push([])
        for(var j=0;j<D;j++){
            x_[i].push(0)
            x_[i][j]=x[i][j]-avg[j]
        }
    }

    //sigma初始值已设定，协方差
    var sigma_ = multiply(transpose(x_), x_, 1/(N-1))
    var sigma__ = multiply(transpose(x_), x_, 1/(N-1))
    var sigma = [sigma_,sigma__]
    var lsigma_ = multiply(transpose(x_), x_, 1/(N-1))
    var lsigma__ = multiply(transpose(x_), x_, 1/(N-1))
    var lsigma = [lsigma_,lsigma__]

    this.mu = mu
    this.sigma = sigma
    this.pi = pi
    this.lmu = lmu
    this.lsigma = lsigma
    this.lpi = lpi

    const modelString = String.raw
    `$$ \begin{aligned} \mu_1 &:=\left( \begin{array}{l} ${mu[0][0]} & ${mu[0][1]} \end{array} \right) \\ \mu_2 &:=\left( \begin{array}{l} ${mu[1][0]} & ${mu[1][1]} \end{array} \right) \\ \sigma^2_1&:= \left( \begin{array}{l} ${sigma[0][0][0]} & ${sigma[0][0][1]}\\ ${sigma[0][1][0]} & ${sigma[0][1][1]} \end{array} \right)   \\ \sigma^2_2&:= \left( \begin{array}{l} ${sigma[1][0][0]} & ${sigma[1][0][1]}\\ ${sigma[1][1][0]} & ${sigma[1][1][1]} \end{array} \right)   \\  \alpha_1 &=${pi[0]} \\ \alpha_2 & =${pi[1]} \end{aligned} $$`

    mdRender(`${modelString}`, this.modelDiv);

    // 显示数据集
    if(this.trainDS){
      ReactDom.unmountComponentAtNode(this.inputDiv);
      ReactDom.render(
        <React.StrictMode>
          <MatView2D 
            Mat     = {new TMatrix.Mat(this.trainDS.Mat)}
            Labels  = {this.trainDS.Labels}
            XRange  = {[-5, 5]}
            YRange  = {[-5, 5]}
            Separator = { {type:'line', w:w.toVector(), b:b} }
            PointStyleDict = {this.pointStyleDict}
          />
        </React.StrictMode>,
        this.inputDiv
      );
    }
  }

  // 输入数据UI
  initInputUI(){

    // 文件选择器
    const fileSelector      = document.createElement('input');
    fileSelector.type       = 'file';
    fileSelector.hidden     = true;
    fileSelector.onchange   = this.handleFileChange.bind(this);   
    
    // this.addButton(this.domElementList.input, "获取服务器数据", this.initTrainDS.bind(this));
    TDom.createButton(this.domElementList.input, "获取服务器数据", this.initTrainDS.bind(this));
    this.addButton(this.domElementList.input, "使用本地数据", this.openFileSelector.bind(this));
    this.domElementList.input.appendChild(fileSelector);
    this.addButton(this.domElementList.input, "训练", this.train.bind(this));


    this.labelTrainTips = this.addLabel(this.domElementList.input, '');
    this.inputDiv       = this.addDiv(this.domElementList.input);

    this.fileSelector  = fileSelector;
  }

  // 输出UI
  initOutputUI(){
    this.outputDiv   = this.addDiv(this.domElementList.output);
    this.trainLogDiv = this.addDiv(this.domElementList.output);
  }

  // 获取描述md文件
  initDescription(){
    TDom.getDataFromUrl(this.descriptionUrl, (txt)=>{
      mdRender(txt, this.descriptionDiv);
    })

  }

  // 获取服务器数据集
  initTrainDS(){
    this.prepare(this.datasetUrl);
  }

  // 准备数据训练数据
  prepare(url){
    TDom.getDataFromUrl(url, 
      (txt)=>{
      this.dataset = txtToMatrix(txt);
      this.trainDS = TCommon.loadDataset(this.dataset);
      this.pointStyleDict = createRandomPointStyleDict(this.trainDS.Labels);
      this.updateModel();
    })
  }

  openFileSelector(){
    this.fileSelector.click();
  }
  // 获取文件路径加载
  handleFileChange(event){
    const files = event.target.files;

    for(let i = 0, il = files.length; i < il; i++){
      let reader = new FileReader();
      reader.onload = (e) => {      
        this.prepare(e.target.result);
        if(i === files.length - 1){
          event.target.value = '';
        }
      };
      reader.readAsDataURL(files[i])
    }
  }

  train(){
    var x = this.x
    let k = 2, N = x.length, D = x[0].length

    //pi初始值已设定
    var pi = this.pi
    var lpi = this.lpi

    var mu = this.mu
    var lmu = this.lmu

    var sigma = this.sigma
    var lsigma = this.lsigma

    this.trainLog = []
    var count = 0
    this.trainLog.push(["迭代次数","均值\mu","协方差\sigma^2","权值系数\alpha"])

    while(true){
      //E - Step
      var psi = []
      for(var i=0;i<k;i++){
          psi.push(guass(mu[i],sigma[i],x))
      }

      var sum_gamma = []
      for(var i=0;i<N;i++){
        sum_gamma.push(0)
        for(var j=0;j<k;j++){
          sum_gamma[i] += psi[j][i] * pi[j]
        }
      }

      var gamma = []
      for(var i=0;i<k;i++){
        gamma.push([])
        for(var j=0;j<N;j++){
          gamma[i].push(pi[i]*psi[i][j]/sum_gamma[j]) 
        }
      }
      if(isNaN(gamma[0][0])){console.log('1234'); break}

      //M - Step
      mu = new Array(lmu.length).fill([]).map(arr => new Array(lmu[0].length).fill(0))
      sigma = new Array(lsigma.length).fill([]).map(arr => new Array(lsigma[0].length).fill([]).map(arr_ => new Array(lsigma[0][0].length).fill(0)))
          
      for(var i=0;i<k;i++){
        //updata mu
        var sum_mu = new Array(mu[i].length).fill(0)
        var gamma_sum = 0
        for(var j=0;j<N;j++){
          for(var ii=0;ii<sum_mu.length;ii++){
            sum_mu[ii] += gamma[i][j] * x[j][ii]
          }
          gamma_sum += gamma[i][j]
        }
        for(var j=0;j<mu[i].length;j++){
          mu[i][j] = sum_mu[j]/gamma_sum
        }

        //updata sigma
        var sigma_sum = new Array(D).fill(0).map(arr => new Array(D).fill(0))
        for(var j=0;j<N;j++){
          for(var ii=0;ii<sigma_sum.length;ii++){
            for(var jj=0;jj<sigma_sum[i].length;jj++){
              sigma_sum[ii][jj] += gamma[i][j] * (x[j][ii]-lmu[i][ii]) * (x[j][jj]-lmu[i][jj])
            }
          }
        }
        for(var ii=0;ii<sigma_sum.length;ii++){
          for(var jj=0;jj<sigma_sum[i].length;jj++){
            sigma[i][ii][jj] = sigma_sum[ii][jj]/gamma_sum
          }
        }
          
          //updata pi
          pi[i] = gamma_sum/N
      }

      //error test
      var err = new Array(k).fill(0)
      for(var i=0;i<mu.length;i++){
        for(var j=0;j<mu[i].length;j++){
          err[i] += Math.abs(lmu[i][j] - mu[i][j])
        }
      }
      for(var i=0;i<sigma.length;i++){
        for(var j=0;j<sigma[i].length;j++){
          for(var kk=0;kk<sigma[i][j].length;kk++){
            err[i] += Math.abs(lsigma[i][j][kk] - sigma[i][j][kk])
          }
        }
      }
      for(var i=0;i<pi.length;i++){
        err[i] += Math.abs(lpi[i] - pi[i])
      }
      
      if(err[0]<1e-16 && err[1]<1e-16) break
      else if(isNaN(mu[0][0]) || isNaN(sigma[0][0][0]) || isNaN(pi[0]) ||isNaN(err[0])){          
        mu = lmu
        sigma = lsigma
        Object.assign(pi , lpi);      
        break
      }
      else if(count>1000) break
      else{
        lmu = new Array(mu.length).fill([]).map(arr => new Array(mu[0].length).fill(0))
        for(var i=0;i<mu.length;i++){
          lmu[i] = mu[i].slice()
        }
          
        lsigma = new Array(sigma.length).fill([]).map(arr => new Array(sigma[0].length).fill([]).map(arr_ => new Array(sigma[0][0].length).fill(0)))
        for(var i=0;i<sigma.length;i++){
          for(var j=0;j<sigma[i].length;j++){
            lsigma[i][j] = sigma[i][j].slice()
          }
        }

        Object.assign(lpi, pi);
      }
      count++
      this.trainLog.push([`${count}`,`\\mu_1 =\\left( \\begin{array}{l} ${mu[0][0].toFixed(2)} & ${mu[0][1].toFixed(2)} \\end{array} \\right) \\\\ \\mu_2 =\\left( \\begin{array}{l} ${mu[1][0].toFixed(2)} & ${mu[1][1].toFixed(2)} \\end{array} \\right)`,
      `\\sigma^2_1= \\left(\\begin{array}{l} ${sigma[0][0][0].toFixed(2)} & ${sigma[0][0][1].toFixed(2)} \\\\ ${sigma[0][1][0].toFixed(2)} & ${sigma[0][1][1].toFixed(2)} \\end{array}\\right) \\\\ \\sigma^2_2= \\left(\\begin{array}{l} ${sigma[1][0][0].toFixed(2)} & ${sigma[1][0][1].toFixed(2)} \\\\ ${sigma[1][1][0].toFixed(2)} & ${sigma[1][1][1].toFixed(2)} \\end{array}\\right)`,
      `\\alpha_1 =${pi[0].toFixed(2)} \\\\ \\alpha_2 =${pi[1].toFixed(2)}`])
    }

    for(var i=0;i<x.length;i++){
      var min = psi[0][i] * pi[0]
      for(var j=1;j<psi.length;j++){
        if(psi[j][i]< min){
          min = psi[j][i] * pi[j]
          this.trainDS.Labels[i] = String(j)
        }
      }
    }

    var w = new Mat.zeros([1, 2]);
    var b = 0.0;
    if(this.trainDS){
      ReactDom.unmountComponentAtNode(this.inputDiv);
      ReactDom.render(
        <React.StrictMode>
          <MatView2D 
            Mat     = {new TMatrix.Mat(this.trainDS.Mat)}
            Labels  = {this.trainDS.Labels}
            XRange  = {[-5, 5]}
            YRange  = {[-5, 5]}
            Separator = { {type:'line', w:w.toVector(), b:b} }
            PointStyleDict = {this.pointStyleDict}
          />
        </React.StrictMode>,
        this.inputDiv
      );
    }
    
    if(this.testDiv) this.testDiv.innerHTML = '';
    if(this.trainLogDiv) this.testDiv.innerHTML = '';
    const modelString = String.raw
    `$$ \begin{aligned} \mu_1 &:=\left( \begin{array}{l} ${mu[0][0]} & ${mu[0][1]} \end{array} \right) \\ \mu_2 &:=\left( \begin{array}{l} ${mu[1][0]} & ${mu[1][1]} \end{array} \right) \\ \sigma^2_1&:= \left( \begin{array}{l} ${sigma[0][0][0]} & ${sigma[0][0][1]}\\ ${sigma[0][1][0]} & ${sigma[0][1][1]} \end{array} \right)   \\ \sigma^2_2&:= \left( \begin{array}{l} ${sigma[1][0][0]} & ${sigma[1][0][1]}\\ ${sigma[1][1][0]} & ${sigma[1][1][1]} \end{array} \right)   \\  \alpha_1 &=${pi[0]} \\ \alpha_2 & =${pi[1]} \end{aligned} $$`
    mdRender(`${modelString}`, this.trainLogDiv);
    showMathStringInTable(this.testDiv, this.trainLog);
  }

  
}


export default EMForEoPOfGMM;